#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing.h>
#include <dlib/image_io.h>
#include <dlib/dnn.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <sstream>
#include <resnet.h>

using namespace dlib;
using namespace std;

anet_type net;

struct FaceRecord {
    string name;
    matrix<float,0,1> descriptor;
};

std::vector<FaceRecord> load_face_lib(const string& filename) {
    std::vector<FaceRecord> face_lib;
    ifstream file(filename);
    string line;
    
    while (getline(file, line)) {
        istringstream iss(line);
        FaceRecord record;
        string value;
        
        getline(iss, record.name, ',');
        
        matrix<float,0,1> descriptor(128);
        for (long i = 0; i < 128; ++i) {
            getline(iss, value, ',');
            descriptor(i) = stof(value);
        }
        record.descriptor = descriptor;
        face_lib.push_back(record);
    }
    
    return face_lib;
}

string recognize_face(const matrix<float,0,1>& descriptor,
                     const std::vector<FaceRecord>& face_lib,
                     double threshold = 0.6) {
    string best_match = "未知";
    double min_distance = 1.0;
    
    for (const auto& record : face_lib) {
        double distance = length(descriptor - record.descriptor);
        if (distance < min_distance) {
            min_distance = distance;
            best_match = record.name;
        }
    }
    
    if (min_distance > threshold) {
        return "陌生人";
    }
    
    return best_match;
}

int main(int argc, char** argv) {
    if (argc != 2) {
        cout << "用法: ./face_recognition </home/hhy212306170/practice/dlib/myapp/2.jpg>" << endl;
        return -1;
    }

    std::vector<FaceRecord> face_lib = load_face_lib("/home/hhy212306170/practice/dlib/myapp/build/facelib.csv");
    if (face_lib.empty()) {
        cerr << "人脸数据库加载失败或为空" << endl;
        return -1;
    }

    frontal_face_detector detector = get_frontal_face_detector();

    shape_predictor sp;
    deserialize("/home/hhy212306170/practice/model/shape_predictor_5_face_landmarks.dat") >> sp;

    deserialize("/home/hhy212306170/practice/model/dlib_face_recognition_resnet_model_v1.dat") >> net;

    matrix<rgb_pixel> img;
    load_image(img, argv[1]);

    std::vector<rectangle> faces = detector(img);
    cout << "检测到 " << faces.size() << " 张人脸，开始识别..." << endl;

    for (size_t i = 0; i < faces.size(); ++i) {

        auto shape = sp(img, faces[i]);
        
        matrix<rgb_pixel> face_chip;
        extract_image_chip(img, get_face_chip_details(shape, 150, 0.25), face_chip);

        matrix<float,0,1> face_descriptor = net(face_chip);
        string name = recognize_face(face_descriptor, face_lib);
        cout << "人脸 " << i+1 << ": " << name << endl;
    }

    return 0;
}